#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
DM_VERITY handler class, supports partition-level integrity protection
"""

import os
import subprocess
import tempfile
import shutil
import logging
from typing import Dict, Optional, List, Tuple
import json
import time
import re

logger = logging.getLogger(__name__)

class DMVerityError(Exception):
    """Error during DM_VERITY operations"""
    pass

class DMVerity:
    """DM_VERITY handler class"""
    
    def __init__(self):
        """Initialize"""
        self._check_dependencies()
    
    def _check_dependencies(self):
        """Check if required tools are installed"""
        required_tools = [
            'veritysetup',
            'cryptsetup',
            'qemu-nbd',
            'parted',
            'mkfs.ext4',
            'blkid',
            'mount',
            'umount'
        ]
        
        for tool in required_tools:
            try:
                subprocess.run(['which', tool], 
                             stdout=subprocess.PIPE, 
                             stderr=subprocess.PIPE, 
                             check=True)
            except (subprocess.SubprocessError, FileNotFoundError):
                raise DMVerityError(f"{tool} tool not found")
    
    def _mount_qcow2(self, image_path: str) -> str:
        """
        Mount qcow2 image as block device
        
        Args:
            image_path: qcow2 image path
        
        Returns:
            str: nbd device path
        """
        try:
            # Load nbd module
            subprocess.run(['modprobe', 'nbd', 'max_part=16'], check=True)
            
            # Randomly get an nbd device, max 20 times
            for i in range(20):
                import random
                nbd_dev = f'/dev/nbd{random.randint(0, 15)}'
                try:
                    # Check if device is in use
                    subprocess.run(['qemu-nbd','--cache=none', '-c', nbd_dev, image_path], check=True)
                    logger.info(f"Image mounted to {nbd_dev}")
                    return nbd_dev
                except subprocess.CalledProcessError:
                    continue
                    
            raise DMVerityError("No available nbd device found")
            
        except subprocess.CalledProcessError as e:
            raise DMVerityError(f"Failed to mount qcow2 image: {e}")
    
    def _unmount_qcow2(self, nbd_dev: str):
        """
        Unmount nbd device
        
        Args:
            nbd_dev: nbd device path
        """
        try:
            subprocess.run(['qemu-nbd', '-d', nbd_dev], check=True)
            logger.info(f"Unmounted {nbd_dev}")
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to unmount nbd device: {e}")
    
    def _find_rootfs_partition(self, nbd_dev: str) -> str:
        """
        Find rootfs partition
        
        Args:
            nbd_dev: nbd device path
        
        Returns:
            str: rootfs partition device path
        """
        try:
            # Use blkid to find all partitions
            import glob
            part_pattern = f"{nbd_dev}p*"
            part_devs = sorted(glob.glob(part_pattern))
            if not part_devs:
                raise DMVerityError(f"{nbd_dev} 未找到任何分区")
            # blkid支持通配符，直接传递所有分区
            result = subprocess.run(['blkid'] + part_devs,
                                    stdout=subprocess.PIPE,
                                    stderr=subprocess.PIPE,
                                    text=True,
                                    check=True)
            
            # Analyze output to find rootfs partition
            for line in result.stdout.splitlines():
                if nbd_dev in line and ('LABEL="cloudimg-rootfs"' in line or 'LABEL="/"' in line):
                    return line.split(':')[0]
            
            # If not found by label, try to find the largest partition
            result = subprocess.run(['parted', '-s', nbd_dev, 'print'],
                                 stdout=subprocess.PIPE,
                                 stderr=subprocess.PIPE,
                                 text=True,
                                 check=True)
            
            max_size = 0
            root_part = None
            for line in result.stdout.splitlines():
                if 'ext4' in line:
                    parts = line.split()
                    if len(parts) >= 4:
                        size = float(parts[3].replace('GB','').replace('MB',''))
                        if size > max_size:
                            max_size = size
                            root_part = f"{nbd_dev}p{parts[0]}"
            
            if root_part:
                return root_part
                
            raise DMVerityError("rootfs partition not found")
            
        except subprocess.CalledProcessError as e:
            raise DMVerityError(f"Failed to find rootfs partition: {e}")
    

    
    def _setup_initramfs_hook(self, image_path: str):
        """
        Set initramfs hook script
        
        Args:
            image_path: image file path
        """
        # Check if virt-customize is installed
        try:
            subprocess.run(['which', 'virt-customize'], check=True)
        except subprocess.CalledProcessError:
            raise DMVerityError("virt-customize tool not found, please install libguestfs-tools")
        
        # Get script files paths
        script_files = self._get_initramfs_script_files()
        
        try:
            # Build virt-customize command
            cmd = [
                'virt-customize',
                '-a', image_path,  # Specify image path
                # Copy hooks directory
                '--upload', f'{script_files["verity_hook"]}:/etc/initramfs-tools/hooks/verity_hook',
                # Copy init-premount directory
                '--upload', f'{script_files["verity_open"]}:/etc/initramfs-tools/scripts/init-premount/verity_open',
                '--upload', f'{script_files["order"]}:/etc/initramfs-tools/scripts/init-premount/ORDER',
                # Set permissions
                '--run-command', 'chmod 755 /etc/initramfs-tools/hooks/verity_hook /etc/initramfs-tools/scripts/init-premount/verity_open',
                # Update initramfs
                '--run-command', 'update-initramfs -u'
            ]
            
            # Execute command
            subprocess.run(cmd, check=True)
            logger.info("Successfully set initramfs hook and updated initrd")
            
        except subprocess.CalledProcessError as e:
            raise DMVerityError(f"Failed to update initrd using virt-customize: {e}")
        finally:
            # Clean up temporary files if created
            self._cleanup_temp_files(script_files)
    
    def _get_initramfs_script_files(self) -> Dict[str, str]:
        """
        Get initramfs script files, using pkg_resources if available, otherwise fallback to filesystem
        
        Returns:
            Dict[str, str]: dict with script file paths
        """
        try:
            import pkg_resources
            
            # Create temporary directory to extract package data files
            temp_dir = tempfile.mkdtemp(prefix='trustimage-initramfs-')
            
            # Extract script files from package
            script_files = {
                'verity_hook': ('script/initramfs/hooks/verity_hook', 'verity_hook'),
                'verity_open': ('script/initramfs/init-prmount/verity_open', 'verity_open'),
                'order': ('script/initramfs/init-prmount/ORDER', 'ORDER')
            }
            
            file_paths = {}
            for key, (resource_path, filename) in script_files.items():
                 try:
                     content = pkg_resources.resource_string('trustimage.core', resource_path)
                     file_path = os.path.join(temp_dir, filename)
                     with open(file_path, 'wb') as f:
                         f.write(content)
                     file_paths[key] = file_path
                 except Exception as e:
                     # If package resource not found, fall back to filesystem
                     logger.warning(f"Failed to load resource {resource_path}: {e}")
                     raise ImportError(f"Resource not found: {resource_path}")
            
            # Mark temp directory for cleanup
            file_paths['_temp_dir'] = temp_dir
            return file_paths
            
        except ImportError:
            # Fallback to original method if pkg_resources is not available
            logger.warning("pkg_resources not available, falling back to filesystem access")
            script_base_path = os.path.join(os.path.dirname(__file__), 'script', 'initramfs')
            
            return {
                'verity_hook': os.path.join(script_base_path, "hooks/verity_hook"),
                'verity_open': os.path.join(script_base_path, "init-prmount/verity_open"),
                'order': os.path.join(script_base_path, "init-prmount/ORDER")
            }
    
    def _cleanup_temp_files(self, script_files: Dict[str, str]):
        """Clean up temporary files if they were created"""
        temp_dir = script_files.get('_temp_dir')
        if temp_dir and os.path.exists(temp_dir):
            shutil.rmtree(temp_dir, ignore_errors=True)
    
    def _align_size(self, size: int, alignment: int = 512) -> int:
        """Align size to specified boundary
        
        Args:
            size: original size (bytes)
            alignment: alignment size (bytes), default 512
        
        Returns:
            int: aligned size (bytes)
        """
        return ((size + alignment - 1) // alignment) * alignment

    def _get_image_virtual_size(self, image_path: str) -> int:
        """Get virtual size (bytes) of image file"""
        try:
            result = self._run_command(
                ['qemu-img', 'info', '--output=json', image_path],
                "Failed to get image info"
            )
            info = json.loads(result)
            return int(info['virtual-size'])
        except (json.JSONDecodeError, KeyError, DMVerityError) as e:
            raise DMVerityError(f"Failed to get virtual size from qemu-img info: {e}")

    def _run_command(self, cmd: List[str], error_msg: str) -> str:
        """Run command and return output
        
        Args:
            cmd: command and argument list
            error_msg: error message prefix
        
        Returns:
            str: command output
        
        Raises:
            DMVerityError: if command fails
        """
        try:
            result = subprocess.run(cmd,
                                 stdout=subprocess.PIPE,
                                 stderr=subprocess.PIPE,
                                 text=True,
                                 check=True)
            return result.stdout
        except subprocess.CalledProcessError as e:
            logger.error(f"{error_msg}: stdout={e.stdout}, stderr={e.stderr}")
            raise DMVerityError(f"{error_msg}: {e}")

    def _get_disk_info(self, device: str) -> Dict:
        """Get disk info
        
        Returns:
            Dict: disk info dict
            {
                'sector_size': int,      # sector size (bytes)
                'total_sectors': int,    # total sectors
                'alignment': int,        # alignment (sectors)
                'last_sector': int       # last used sector
            }
        """
        info = {
            'sector_size': 512,  # default
            'total_sectors': 0,
            'alignment': 2048,   # default 1MB alignment (2048 * 512 = 1MB)
            'last_sector': 0
        }
        
        try:
            # 1. Try blockdev to get basic info
            try:
                # Get sector size
                sector_size = self._run_command(
                    ['blockdev', '--getss', device],
                    "Failed to get sector size"
                )
                info['sector_size'] = int(sector_size.strip())
                logger.debug(f"Sector size: {info['sector_size']}B")
                
                # Get total sectors
                total_sectors = self._run_command(
                    ['blockdev', '--getsz', device],
                    "Failed to get total sectors"
                )
                info['total_sectors'] = int(total_sectors.strip())
                logger.debug(f"Total sectors: {info['total_sectors']}")
                
            except (DMVerityError, ValueError) as e:
                logger.warning(f"Failed to get info using blockdev: {e}")
            
            # 2. Use parted for detailed info
            parted_output = self._run_command(
                ['parted', '-s', device, 'unit', 's', 'print'],
                "Failed to get partition info"
            )
            logger.debug(f"parted output:\n{parted_output}")
            
            # Parse parted output
            for line in parted_output.splitlines():
                line = line.strip()
                
                # If sector size not obtained, parse from parted output
                if info['sector_size'] == 512 and 'Sector size' in line:
                    try:
                        sizes = line.split(':')[1].strip()
                        logical_size = sizes.split('/')[0].strip()
                        info['sector_size'] = int(logical_size.replace('B', ''))
                        logger.debug(f"Got sector size from parted: {info['sector_size']}")
                    except (ValueError, IndexError) as e:
                        logger.warning(f"Failed to parse sector size: {e}, line: {line}")
                
                # If total sectors not obtained, parse from parted output
                elif info['total_sectors'] == 0 and line.startswith('Disk ' + device):
                    try:
                        size_part = line.split(':')[1].strip()
                        size_str = size_part.replace('sectors', '').replace('s', '').strip()
                        info['total_sectors'] = int(size_str)
                        logger.debug(f"Got total sectors from parted: {info['total_sectors']}")
                    except (ValueError, IndexError) as e:
                        logger.warning(f"Failed to parse total sectors: {e}, line: {line}")
                
                # Parse partition table to get last used sector
                elif line and line[0].isdigit():
                    try:
                        parts = line.split()
                        if len(parts) >= 3:
                            end_sector = int(parts[2].rstrip('s'))
                            info['last_sector'] = max(info['last_sector'], end_sector)
                            logger.debug(f"Update last sector position: {info['last_sector']}")
                    except (ValueError, IndexError) as e:
                        logger.warning(f"Failed to parse partition end position: {e}, line: {line}")
            
            # 3. If still can't get total sectors, try fdisk
            if info['total_sectors'] == 0:
                try:
                    fdisk_output = self._run_command(
                        ['fdisk', '-l', device],
                        "Failed to get disk info"
                    )
                    for line in fdisk_output.splitlines():
                        if 'sectors' in line and device in line:
                            try:
                                # Format: Disk /dev/nbd0: 1 GiB, 1073741824 bytes, 2097152 sectors
                                sectors = line.split('sectors')[0].strip().split()[-1]
                                info['total_sectors'] = int(sectors)
                                logger.debug(f"Got total sectors from fdisk: {info['total_sectors']}")
                                break
                            except (ValueError, IndexError) as e:
                                logger.warning(f"Failed to parse fdisk output: {e}, line: {line}")
                except DMVerityError:
                    logger.warning("Failed to get info using fdisk")
            
            # 4. Validate required info
            if info['total_sectors'] == 0:
                raise DMVerityError("Failed to get total disk sectors")
            
            logger.debug(f"Disk info: sector size={info['sector_size']}B, "
                       f"total sectors={info['total_sectors']}, "
                       f"last used sector={info['last_sector']}, "
                       f"total size={info['total_sectors'] * info['sector_size'] / (1024**3):.2f}GB")
            
            return info
            
        except Exception as e:
            logger.error(f"Failed to get disk info: {e}")
            raise DMVerityError(f"Failed to get disk info: {e}")
    
    def _configure_var_mount(self, mount_point: str, var_part: str):
        """Configure var partition mount"""
        try:
            logger.info(f"Formatting var partition: {var_part}")
            self._run_command(['mkfs.ext4', '-L', 'var', var_part],
                               "Failed to format var partition")
            # Get var partition UUID
            var_uuid = subprocess.run(['blkid', '-s', 'UUID', '-o', 'value', var_part],
                                   stdout=subprocess.PIPE,
                                   text=True,
                                   check=True).stdout.strip()
            
            # Configure fstab
            fstab_path = os.path.join(mount_point, 'etc/fstab')
            with open(fstab_path, 'a') as f:
                f.write(f'\n# /var partition\n')
                f.write(f'UUID={var_uuid} /var ext4 rw,noexec,nosuid,nodev 0 2\n')
                # TODO: Copy files from original /var directory to new partition
                
        except (subprocess.CalledProcessError, IOError) as e:
            raise DMVerityError(f"Failed to configure var partition mount: {e}")
    
    def _generate_verity_hash(self, root_part: str, hash_part: str, salt: Optional[str] = None) -> Dict[str, str]:
        """Generate dm-verity hash data and write directly to partition
        
        Args:
            root_part: rootfs partition device path
            hash_part: verity hash partition device path
            salt: optional salt
        
        Returns:
            Dict[str, str]: dict with hash info
        """
        try:
            salt_param = ['--salt', salt] if salt else []
            
            # Use veritysetup format to write hash directly to partition
            cmd = ['veritysetup', 'format', root_part, hash_part] + salt_param
            logger.info(f"veritysetup format command: {cmd}")
            result = subprocess.run(cmd,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE,
                                text=True,
                                check=True)

            # Run sync to ensure data is written to disk
            logger.debug("Run sync to ensure data is written to disk")
            subprocess.run(['sync'], check=True)
            
            # Parse veritysetup output
            return self._parse_verity_output(result.stdout)
            
        except subprocess.CalledProcessError as e:
            raise DMVerityError(f"Failed to generate dm-verity hash: {e}")
            
        except Exception as e:
            raise DMVerityError(f"Failed to generate dm-verity hash: {e}")
    
    def _collect_partition_info(self, root_part: str, var_part: str, 
                              hash_part: str) -> Dict[str, str]:
        """Collect UUID and other info for all partitions"""
        try:
            root_uuid = subprocess.run(['blkid', '-s', 'UUID', '-o', 'value', root_part],
                                    stdout=subprocess.PIPE,
                                    text=True,
                                    check=True).stdout.strip()
            
            var_uuid = subprocess.run(['blkid', '-s', 'UUID', '-o', 'value', var_part],
                                    stdout=subprocess.PIPE,
                                    text=True,
                                    check=True).stdout.strip()

            hash_uuid = subprocess.run(['blkid', '-s', 'UUID', '-o', 'value', hash_part],
                                     stdout=subprocess.PIPE,
                                     text=True,
                                     check=True).stdout.strip()

            return {
                'root_uuid': root_uuid,
                'var_uuid': var_uuid,
                'hash_uuid': hash_uuid
            }
            
        except subprocess.CalledProcessError as e:
            raise DMVerityError(f"Failed to collect partition info: {e}")

    def _update_kernel_params(self, nbd_device: str, root_uuid: str, verity_uuid: str, roothash: str) -> bool:
        """
        Update kernel boot parameters
        
        Args:
            nbd_device: mounted NBD device path
            root_uuid: UUID of root partition
            verity_uuid: UUID of verity partition
            roothash: root hash value
        Returns:
            bool: success or not
        """
        mount_point = None
        
        try:
            # 1. Find boot partition
            boot_partition = None
            for i in range(20):
                part_device = f"{nbd_device}p{i}"
                if not os.path.exists(part_device):
                    continue
                    
                # Create temp mount point to check partition
                temp_mount = tempfile.mkdtemp(prefix='boot-check-')
                try:
                    subprocess.run(['mount', '-o', 'ro', part_device, temp_mount], check=True)
                    if os.path.isdir(os.path.join(temp_mount, 'grub')):
                        boot_partition = part_device
                    subprocess.run(['umount', temp_mount], check=True)
                    os.rmdir(temp_mount)
                    if boot_partition:
                        break
                except:
                    subprocess.run(['umount', temp_mount], check=False)
                    os.rmdir(temp_mount)
                    continue
                    
            if not boot_partition:
                raise Exception("boot partition not found")
                
            # 2. Mount boot partition
            mount_point = tempfile.mkdtemp(prefix='boot-mount-')
            subprocess.run(['mount', boot_partition, mount_point], check=True)
            
            # 3. Read and modify grub.cfg
            grub_cfg_path = os.path.join(mount_point, 'grub/grub.cfg')
            if not os.path.isfile(grub_cfg_path):
                raise Exception(f"{grub_cfg_path} not found")
                
            # Backup original file
            shutil.copy2(grub_cfg_path, f"{grub_cfg_path}.bak")
            
            # Read file content
            with open(grub_cfg_path, 'r') as f:
                content = f.read()
                
            # Build new parameters
            new_params = {
                'rootuuid': f"{root_uuid}",
                'verityuuid': f"{verity_uuid}",
                'roothash': roothash
            }
            
            def parse_kernel_params(cmdline: str) -> dict:
                params = {}
                parts = cmdline.split()
                for part in parts:
                    if '=' in part:
                        key, value = part.split('=', 1)
                        params[key] = value
                    else:
                        params[part] = ""
                return params
                
            def build_kernel_params(params: dict) -> str:
                parts = []
                # Ensure replace root param with /dev/mapper/root
                params['root'] = '/dev/mapper/root'
                for key, value in params.items():
                    if value:
                        parts.append(f"{key}={value}")
                    else:
                        parts.append(key)
                return ' '.join(parts)
                
            def update_cmdline(match):
                line = match.group(0)
                if '/vmlinuz' not in line:  # Check if contains vmlinuz
                    return line
                    
                                # Find full vmlinuz path
                try:
                    # Find full vmlinuz path
                    vmlinuz_start = line.find('/vmlinuz')
                    if vmlinuz_start == -1:
                        return line
                        
                    # Get part before vmlinuz
                    remaining = line[vmlinuz_start:]
                    vmlinuz_parts = remaining.split(None, 1)
                    if len(vmlinuz_parts) != 2:
                        return line
                        
                    vmlinuz_path = vmlinuz_parts[0]  # Full vmlinuz path
                    cmdline = vmlinuz_parts[1]  # Remaining command line arguments
                    
                    # Get part before vmlinuz
                    prefix = line[:vmlinuz_start]
                    
                    current_params = parse_kernel_params(cmdline)
                    current_params.update(new_params)
                    new_cmdline = build_kernel_params(current_params)
                    
                    return f"{prefix}{vmlinuz_path} {new_cmdline}"
                except Exception as e:
                    logger.error(f"Error processing kernel command line: {e}")
                    return line
                
            # Update kernel params
            pattern = r'linux\s+[^\n]+'
            new_content = re.sub(pattern, update_cmdline, content)
            
            logger.debug(f"Updated grub.cfg content: {new_content}")
            # Write back to file
            with open(grub_cfg_path, 'w') as f:
                f.write(new_content)
                
            # Sync filesystem
            subprocess.run(['sync'], check=True)
            
            return True
            
        except Exception as e:
            logger.error(f"Failed to update kernel params: {e}")
            return False
            
        finally:
            # Clean up resources
            try:
                if mount_point:
                    if os.path.ismount(mount_point):
                        subprocess.run(['umount', mount_point], check=False)
                    os.rmdir(mount_point)
            except Exception as e:
                logger.error(f"Failed to clean up resources: {e}")

    def _create_partition_and_wait(self, nbd_dev: str, part_name: str, fs_type: str, part_size_bytes: int, timeout_seconds: int = 15) -> str:
        """
        Create partition and wait for device to be ready (auto-calculate start/end and unit)
        Args:
            nbd_dev: nbd device path
            part_name: partition name
            fs_type: filesystem type
            part_size_bytes: partition size (bytes)
            timeout_seconds: timeout for device to be ready (seconds)
        Returns:
            str: real path of partition device
        Raises:
            RuntimeError: timeout waiting for partition device
        """
        # Get start position and unit of last Free Space
        last_end = self._get_last_free_space_start(nbd_dev)
        logger.debug(f"Create {part_name} partition, start: {last_end}")
        # Support MB/GB unit
        if last_end.endswith('GB'):
            start_val = float(last_end.replace('GB',''))
            size_gb = part_size_bytes / (1024**3)
            end_val = start_val + size_gb
            part_end = f"{end_val}GB"
        elif last_end.endswith('MB'):
            start_val = float(last_end.replace('MB',''))
            size_mb = part_size_bytes / (1024**2)
            end_val = start_val + size_mb
            part_end = f"{end_val}MB"
        else:
            raise DMVerityError(f"Unrecognized partition unit: {last_end}")
        logger.debug(f"Create {part_name} partition: {last_end} -> {part_end}")
        self._run_command(['parted', '-s', nbd_dev, 'mkpart', part_name, fs_type, last_end, part_end], f"Failed to create {part_name} partition")
       
        part_number= self.find_partition_number_by_label(nbd_dev,part_name)
        if part_number != -1:
            part_realpath=f"{nbd_dev}p{part_number}"
        else:
            raise DMVerityError("Partition number not found")
        logger.debug(f"Successfully found {part_name} partition by name: {part_realpath}")

        # TODO: wait for partition to be ready
        time.sleep(1)
        return part_realpath

    def setup(self, image_path: str, fs_type: str = 'rootfs', salt: Optional[str] = None) -> Dict[str, str]:
        """
        Set up DM_VERITY protection for image
        
        Optimized flow:
        1. Set up initramfs hook (modify rootfs)
        2. Calculate space and resize (ensure 512-byte alignment)
        3. Create required partitions and generate hash
        
        Args:
            image_path: image path
            fs_type: filesystem type, usually 'rootfs'
            salt: salt, if None will be randomly generated
        
        Returns:
            Dict[str, str]: dict with DM_VERITY setup info
        """
        if not os.path.exists(image_path):
            raise DMVerityError(f"Image file does not exist: {image_path}")
        
        logger.info(f"Set up DM_VERITY protection for image {image_path}")
        
        try:
            # 1. Set up initramfs hook (this modifies rootfs content)
            logger.info("Set up initramfs hook")
            self._setup_initramfs_hook(image_path)
            
            # 2. Calculate space requirements and resize
            # 2.1 Get original image virtual size (bytes)
            logger.info(f"Start resizing original image virtual size")
            original_size = self._align_size(self._get_image_virtual_size(image_path))
            logger.info(f"Original image virtual size: {original_size / (1024**3):.2f}GB")

            # 2.2 Set fixed partition sizes
            var_size = 2 * 1024**3       # 2GB for /var
            verity_size = 500 * 1024**2  # 500MB for dm-verity
            buffer_size = 100 * 1024**2  # 100MB buffer
            
            # 2.3 Calculate total size and resize image
            total_size = original_size + var_size + verity_size + buffer_size
            logger.info(f"Resize image to: {total_size / (1024**3):.2f}GB")
            self._run_command(['qemu-img', 'resize', image_path, str(total_size)], "Failed to resize image")
            
            # 3. Mount image and handle partitions
            nbd_dev = None
            try:
                # 3.1 Mount image
                nbd_dev = self._mount_qcow2(image_path)
                logger.debug(f"Image mounted to: {nbd_dev}")
                
                # 3.2 Fix/refresh partition table
                self._run_command(['sgdisk', '-e', nbd_dev ],  "Failed to fix partition table")
                self._run_command(['partprobe', nbd_dev], "Failed to refresh partition table")
                time.sleep(1)  # Wait for partition table update
                   
                # 3.3 Create var partition
                logger.info(f"Create var partition")
                var_part = self._create_partition_and_wait(nbd_dev, "var", "ext4", var_size)
                logger.info(f"var partition created: {var_part}")

                # 3.4 Mount rootfs and configure var partition
                logger.info(f"Configure var partition...")
                root_part = self._find_rootfs_partition(nbd_dev)
                logger.info(f"Found rootfs partition: {root_part}")

                mount_point = tempfile.mkdtemp(prefix='rootfs-')
                try:
                    self._run_command(['mount', root_part, mount_point],
                                   f"Failed to mount {root_part}")
                    self._configure_var_mount(mount_point, var_part)
                finally:
                    if os.path.ismount(mount_point):
                        self._run_command(['umount', mount_point],
                                       f"Failed to unmount {mount_point}")
                    shutil.rmtree(mount_point, ignore_errors=True)
                logger.info(f"var partition configured")

                # 3.4 Create dm-verity partition
                logger.info(f"Create dm-verity partition...")
                hash_part = self._create_partition_and_wait(nbd_dev, "verity-hash", "ext4", verity_size)
                logger.info(f"dm-verity partition created: {hash_part}")

                # 3.5 Generate dm-verity hash and write directly to partition
                logger.info(f"Configure dm-verity...")
                verity_info = self._generate_verity_hash(root_part, hash_part, salt)
                logger.info(f"dm-verity configured: verity_info: {verity_info}")

                # 3.6 Collect partition info
                partition_info = self._collect_partition_info(root_part, var_part, hash_part)
                logger.debug(f"partition info: {partition_info}")
          
                # 3.7 Update kernel params
                logger.info("Update kernel boot parameters...")
                if not self._update_kernel_params(
                    nbd_device=nbd_dev,
                    root_uuid=partition_info['root_uuid'],
                    verity_uuid=partition_info['hash_uuid'],
                    roothash=verity_info['Root hash']
                ):
                    raise DMVerityError("Failed to update kernel boot parameters")
                    
                # 8. Return all info
                return {
                    **verity_info,
                    **partition_info,
                    'fs_type': fs_type,
                    'hash_offset': '0'
                }
                
            finally:
                if nbd_dev:
                    self._unmount_qcow2(nbd_dev)
                    
        except Exception as e:
            logger.error(f"Failed to set up DM_VERITY protection: {e}", exc_info=True)
            raise DMVerityError(f"Failed to set up DM_VERITY protection: {e}")
    
    def _parse_verity_output(self, output: str) -> Dict[str, str]:
        """
        Parse veritysetup output
        
        Args:
            output: output of veritysetup command
        
        Returns:
            Dict[str, str]: parsed DM_VERITY info
        """
        info = {}
        for line in output.splitlines():
            if ':' in line:
                key, value = line.split(':', 1)
                info[key.strip()] = value.strip()
        return info

    def _get_last_free_space_start(self, nbd_dev: str) -> str:
        """
        Get start position of last Free Space section in parted output
        Args:
            nbd_dev: nbd device path
        Returns:
            str: start position of last Free Space section
        Raises:
            DMVerityError: if unable to determine partition start position
        """
        parted_output = self._run_command(['parted', '-s', nbd_dev, 'print', 'free'],
                                       "Failed to get partition layout")
        last_end = None
        for line in parted_output.splitlines():
            if 'Free Space' in line:
                parts = line.split()
                if len(parts) >= 3:
                    last_end = parts[0]
        if not last_end:
            raise DMVerityError("Unable to determine partition start position")
        return last_end

    def find_partition_number_by_label(self, nbd_dev: str, label: str) -> int:
        """
        Find partition number by label
        
        Args:
            device: device path, e.g. /dev/nbd3
            label: partition label, e.g. 'verity-hash'
        
        Returns:
            int: partition number (e.g. 3), -1 if not found
        """
        try:
            # Get partition info
            logger.debug(f"label is {label}")
            result = subprocess.run(['parted', '-s', nbd_dev, 'print'],
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE,
                                text=True,
                                check=True)
            lines = result.stdout.splitlines()
            found_header = False
            for line in lines:
                # Skip content before header
                if not found_header:
                    if line.strip().startswith("Number"):
                        found_header = True
                    continue
                # Skip empty lines
                if not line.strip():
                    continue
                # If line contains label
                if label in line:
                    parts = re.split(r'\s+', line.strip(), maxsplit=6)
                    if len(parts) >= 1:
                        part_number = parts[0]
                        return int(part_number)
            return -1
        except Exception as e:
            print(f"Failed to find partition number: {e}")
            return -1